"""
Created on  2024/10/25
@author: <LIN>
gradio langchain的多轮对话测试。
"""
import gradio as gr

from customize.get_ollama import GetOllama
from langchain.schema import HumanMessage, AIMessage


models = [GetOllama(model_type=1)()]

block_css = """.importantButton {
    background: linear-gradient(45deg, #7e0570,#5d1c99, #6e00ff) !important;
    border: none !important;
}
.importantButton:hover {
    background: linear-gradient(45deg, #ff00e0,#8500ff, #6e00ff) !important;
    border: none !important;
}"""

default_theme_args = dict(
    font=["Source Sans Pro", 'ui-sans-serif', 'system-ui', 'sans-serif'],
    font_mono=['IBM Plex Mono', 'ui-monospace', 'Consolas', 'monospace'],
)

init_message = f"""欢迎使用 ChatGPT Gradio UI！"""


def respond(query, model, temperature, history_turns, chat_history):
    global chat_turns
    llm = models[0]

    history = []
    len_history = min(chat_turns, history_turns)
    if chat_turns > 0:
        for turn in range(len_history):
            history.append(HumanMessage(content=chat_history[len_history - turn][0]));
            history.append(AIMessage(content=chat_history[len_history - turn][1]));

    history.append(HumanMessage(content=query));
    # print(history)

    response = llm.invoke(history).content;

    chat_history.append((query, response));
    chat_turns += 1
    return "", chat_history


def clear(chat_history):
    global chat_turns
    chat_history = [(None, "已清除对话历史")]
    chat_turns = 0
    return chat_history


def setting_change(model, temperature, history_turns, chat_history):
    global chat_turns
    chat_history = [(None, f"设置更新：\n 模型名称：{model} \n 温度：{temperature} \n 记忆历史对话轮数：{history_turns}\n")]
    chat_turns = 0
    return chat_history


with gr.Blocks(css=block_css, theme=gr.themes.Default(**default_theme_args)) as demo:
    gr.Markdown('ChatGPT Gradio')
    chat_turns = 0
    with gr.Row():
        with gr.Column(scale=10):
            chatbot = gr.Chatbot([[None, init_message]],
                                 elem_id="chat-box",
                                 label="聊天历史")
            query = gr.Textbox(label="输入问题",
                               placeholder="请输入提问内容，按回车进行提交")
            clear_button = gr.Button("重新对话", visible=True)

        with gr.Column(scale=5):
            model = gr.Radio(models,
                             label="请选择使用模型",
                             value="models", interactive=True)

            temperature = gr.Slider(0, 1,
                                    value=0.8,
                                    step=0.1,
                                    label="Temperature",
                                    interactive=True)

            history_turns = gr.Slider(1, 20,
                                      value=5,
                                      step=1,
                                      label="对话轮数",
                                      info='记录历史对话轮数',
                                      interactive=True)

            settings_button = gr.Button("更新设置", visible=True)

            settings_button.click(fn=setting_change, inputs=[model, temperature, history_turns, chatbot],
                                  outputs=[chatbot])

    query.submit(respond, [query, model, temperature, history_turns, chatbot], [query, chatbot])

    clear_button.click(fn=clear,
                       inputs=[chatbot],
                       outputs=[chatbot])
    def auto_scroll():
        js_code = """
        () => {
            let chatbox = document.getElementById('chat-box');
            chatbox.scrollTop = chatbox.scrollHeight;
        }
        """
        return gr.update(value=None, elem_id="chat-box", js=js_code)
    chatbot.change(auto_scroll, [], [], every=True)

demo.launch()
